import requests
import numpy as np
from pgvector.django import CosineDistance,L2Distance

from file.models import TextVector, ImageVector

OLLAMA_API = "http://localhost:11434/api/generate"
MODEL_NAME = "qwen2.5vl:7b"

from core.model_loader import get_text_embedding_model

def retrieve_text_vectors(know_base_ids, qusetion, top_k=5):
    # 取出前 5 个最相关的文本向量
    question_vector = get_embedding(qusetion)
    vectors = TextVector.objects.filter(know_base_id__in=know_base_ids)
    results = vectors.annotate(
        distance=L2Distance('embedding_vector', question_vector)
    ).order_by('distance')[:top_k]
    return [result.embedding_content for result in results]

def retrieve_image_vectors(know_base_ids):
    vectors = ImageVector.objects.filter(know_base_id__in=know_base_ids)
    return [v.image_path for v in vectors[:5]]

def generate_with_ollama(prompt):
    payload = {
        "model": MODEL_NAME,
        "prompt": prompt,
        "stream": True
    }

    with requests.post(OLLAMA_API, json=payload, stream=True) as r:
        for line in r.iter_lines():
            if line:
                yield line.decode('utf-8')

def get_embedding(question):
    text_embedding_model = get_text_embedding_model()
    question_embedding = text_embedding_model.get_text_embedding(question)
    return question_embedding